package main

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"

	"golang.org/x/net/dns/dnsmessage"
)

func main() {

	// 启动 HTTPS 服务器，支持动态证书选择
	server := &http.Server{
		Addr: ":443",
		TLSConfig: &tls.Config{
			GetCertificate: getCertificateForSNI,
		},
	}

	// 注册 DoH 请求处理函数
	http.HandleFunc("/dns-query", handleDoHRequest)

	// 启动 HTTPS 服务器
	log.Println("Starting DoH server on https://localhost:443")
	if err := server.ListenAndServeTLS("", ""); err != nil {
		log.Fatal("ListenAndServeTLS: ", err)
	}
}

func getCertificateForSNI(helloInfo *tls.ClientHelloInfo) (*tls.Certificate, error) {
    switch helloInfo.ServerName {
    case "server1.com":
        cert, err := tls.LoadX509KeyPair("server1.crt", "server1.key")
        return &cert, err  // 返回结构体的指针
    case "server2.com":
        cert, err := tls.LoadX509KeyPair("server2.crt", "server2.key")
        return &cert, err  // 返回结构体的指针
    default:
        return nil, fmt.Errorf("no certificate found for SNI: %s", helloInfo.ServerName)
    }
}

// 处理 DoH 请求
func handleDoHRequest(w http.ResponseWriter, req *http.Request) {
	// 读取 DoH 请求中的 DNS 消息
	dnsRequest := req.Body
	defer dnsRequest.Close()

	// 读取请求中的 DNS 消息内容
	var buf bytes.Buffer
	_, err := io.Copy(&buf, dnsRequest)
	if err != nil {
		log.Println("Failed to read DNS request: ", err)
		http.Error(w, "Failed to read DNS request", http.StatusInternalServerError)
		return
	}

	// 解包 DNS 消息
	var msg dnsmessage.Message
	err = msg.Unpack(buf.Bytes())
	if err != nil {
		log.Println("Failed to unpack DNS request: ", err)
		http.Error(w, "Failed to unpack DNS request", http.StatusInternalServerError)
		return
	}

	// 向真实的 DNS 服务器发送请求
	dnsResponse, err := sendDNSRequestToRealServer(msg)
	if err != nil {
		log.Println("Failed to send DNS request to real server: ", err)
		http.Error(w, "Failed to get DNS response", http.StatusInternalServerError)
		return
	}

	// 发送 DNS 响应作为 DoH 响应
	w.Header().Set("Content-Type", "application/dns-message")
	w.Write(dnsResponse)
}

// 向真实的 DNS 服务器发送 DNS 请求
func sendDNSRequestToRealServer(msg dnsmessage.Message) ([]byte, error) {
	// 连接到 DNS 服务器
	dnsServer := "114.114.114.114:53"
	conn, err := net.Dial("udp", dnsServer)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to DNS server: %w", err)
	}
	defer conn.Close()

	// 将 DNS 消息序列化
	dnsBytes, err := msg.Pack()
	if err != nil {
		return nil, fmt.Errorf("failed to pack DNS message: %w", err)
	}

	// 发送 DNS 请求
	_, err = conn.Write(dnsBytes)
	if err != nil {
		return nil, fmt.Errorf("failed to send DNS request: %w", err)
	}

	// 读取 DNS 响应
	responseBuf := make([]byte, 512) // 假设响应不超过 512 字节
	n, err := conn.Read(responseBuf)
	if err != nil {
		return nil, fmt.Errorf("failed to read DNS response: %w", err)
	}

	return responseBuf[:n], nil
}
